# 航向跟踪训练脚本
# TD + Q-learning + ReplayBuffer + GredClip + Double + Dueling + PER + N-step

import gymnasium as gym
import math
import random
from typing import Dict, List, Tuple

import matplotlib.pyplot as plt
from collections import namedtuple, deque
from itertools import count
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

from segment_tree import MinSegmentTree, SumSegmentTree

import socket
import time
import json

import pandas as pd
import pickle

import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

import os
HOME = os.getcwd()
print(HOME)


# if GPU is to be used
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


Transition = namedtuple('Transition',
                        ('state', 'action','reward', 'next_state', 'done'))

class ReplayMemory(object):

    def __init__(self, capacity):
        self.memory = deque([], maxlen=capacity)
        self.max_size = capacity

        # for N-step Learning
        self.n_step_buffer = deque(maxlen=N_STEP)
        self.n_step = N_STEP
        self.gamma = GAMMA

    def push(self, *args):
        """Save a transition"""

        # N-STEP
        self.n_step_buffer.append(Transition(*args))
        if len(self.n_step_buffer) < self.n_step:
            return ()
        # make a n-step transition
        rew, next_obs, done = self._get_n_step_info(
            self.n_step_buffer, self.gamma
        )
        obs = self.n_step_buffer[0].state
        act = self.n_step_buffer[0].action

        self.memory.append(Transition(obs, act, rew, next_obs, done))

        return self.n_step_buffer[0]

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

    def sample_batch_from_idxs(
        self, idxs: np.ndarray
    ) -> Dict[str, np.ndarray]:
        # for N-step Learning
        return [self.memory[index] for index in idxs]
    
    def _get_n_step_info(
        self, n_step_buffer: deque, gamma: float
    ) -> Tuple[np.int64, np.ndarray, bool]:
        """Return n step rew, next_obs, and done."""
        # info of the last transition
        rew = n_step_buffer[-1].reward
        next_obs = n_step_buffer[-1].next_state
        done = n_step_buffer[-1].done

        for transition in reversed(list(n_step_buffer)[:-1]):
            r = transition.reward
            n_o = transition.next_state
            d = transition.done

            rew = r + gamma * rew * (1 - d)
            next_obs, done = (n_o, d) if d else (next_obs, done)

        return rew, next_obs, done
    
    def __len__(self):
        return len(self.memory)


class PrioritizedReplayBuffer(ReplayMemory):
    """Prioritized Replay buffer.

    Attributes:
        max_priority (float): max priority
        tree_ptr (int): next index of tree
        alpha (float): alpha parameter for prioritized replay buffer
        sum_tree (SumSegmentTree): sum tree for prior
        min_tree (MinSegmentTree): min tree for min prior to get max weight

    """

    def __init__(
            self,
            capacity,
            alpha: float = 0.6
    ):
        """Initialization."""
        assert alpha >= 0

        super(PrioritizedReplayBuffer, self).__init__(capacity)
        self.max_priority, self.tree_ptr = 1.0, 0
        self.alpha = alpha

        # capacity must be positive and a power of 2.
        tree_capacity = 1
        while tree_capacity < self.max_size:
            tree_capacity *= 2

        self.sum_tree = SumSegmentTree(tree_capacity)
        self.min_tree = MinSegmentTree(tree_capacity)

    def push(
            self,
            *args
    ):
        """Store experience and priority."""
        transition = super().push(*args)

        if transition:
            self.sum_tree[self.tree_ptr] = self.max_priority ** self.alpha
            self.min_tree[self.tree_ptr] = self.max_priority ** self.alpha
            self.tree_ptr = (self.tree_ptr + 1) % self.max_size

        return transition

    def sample(self, batch_size, beta: float):
        """Sample a batch of experiences."""
        assert len(self) >= batch_size
        assert beta > 0

        indices = self._sample_proportional(batch_size)

        weights = np.array([self._calculate_weight(i, beta) for i in indices])

        return [self.memory[index] for index in indices], weights, indices

    def update_priorities(self, indices: List[int], priorities: np.ndarray):
        """Update priorities of sampled transitions."""
        assert len(indices) == len(priorities)

        for idx, priority in zip(indices, priorities):
            assert priority > 0
            assert 0 <= idx < len(self)

            self.sum_tree[idx] = priority ** self.alpha
            self.min_tree[idx] = priority ** self.alpha

            self.max_priority = max(self.max_priority, priority)

    def _sample_proportional(self, batch_size) -> List[int]:
        """Sample indices based on proportions."""
        indices = []
        p_total = self.sum_tree.sum(0, len(self) - 1)
        segment = p_total / batch_size

        for i in range(batch_size):
            a = segment * i
            b = segment * (i + 1)
            upperbound = random.uniform(a, b)
            idx = self.sum_tree.retrieve(upperbound)
            indices.append(idx)

        return indices

    def _calculate_weight(self, idx: int, beta: float):
        """Calculate the weight of the experience at idx."""
        # get max weight
        p_min = self.min_tree.min() / self.sum_tree.sum()
        max_weight = (p_min * len(self)) ** (-beta)

        # calculate weights
        p_sample = self.sum_tree[idx] / self.sum_tree.sum()
        weight = (p_sample * len(self)) ** (-beta)
        weight = weight / max_weight

        return weight

    
# Deuling Network
class DQN(nn.Module):

    def __init__(self, n_observations: int, n_actions: int):
        super(DQN, self).__init__()

        self.n_observations = n_observations
        self.n_actions = n_actions
        self.features_layer = nn.Linear(n_observations, CELL_NUM) # feature layer
        self.advantage_hidden_layer = nn.Linear(CELL_NUM, CELL_NUM) # advantages layer 1
        self.advantage_layer = nn.Linear(CELL_NUM, n_actions) #  advantages layer 2
        self.value_hidden_layer = nn.Linear(CELL_NUM, CELL_NUM)
        self.value_layer = nn.Linear(CELL_NUM, 1) # value layer 2

    # Called with either one element to determine next action, or a batch
    # during optimization. Returns tensor([[left0exp,right0exp]...]).
    def forward(self, x):
        # Dueling network

        # test for tanh and relu
        x = F.tanh(self.features_layer(x))
        advantage_x = F.tanh(self.advantage_hidden_layer(x))
        advantage = self.advantage_layer(advantage_x)
        value_x = F.tanh(self.value_hidden_layer(x))
        value = self.value_layer(value_x)
        q = value + advantage - advantage.mean(dim=-1, keepdim=True)
        return q

def select_action(state):
    global steps_done
    sample = random.random()
    eps_threshold = EPS_END + (EPS_START - EPS_END) * \
        math.exp(-1. * steps_done / EPS_DECAY)
    steps_done += 1
    if sample > eps_threshold:
        with torch.no_grad():
            # t.max(1) will return the largest column value of each row.
            # second column on max result is index of where max element was
            # found, so we pick action with the larger expected reward.
            return policy_net(state).max(1).indices.view(1, 1), eps_threshold
    else:
        return torch.tensor([[random.choice(action_space)]], device=device, dtype=torch.long), eps_threshold


def plot_durations(show_result=False):
    plt.figure(1)
    durations_t = torch.tensor(episode_durations, dtype=torch.float)
    if show_result:
        plt.title('Result')
    else:
        plt.clf()
        plt.title('Training...')
    plt.xlabel('Episode')
    plt.ylabel('Score')
    plt.plot(durations_t.numpy())
    # Take 100 episode averages and plot them too
    if len(durations_t) >= 100:
        means = durations_t.unfold(0, 100, 1).mean(1).view(-1)
        means = torch.cat((torch.zeros(99), means))
        plt.plot(means.numpy())

def _compute_dqn_loss(transitions, gamma: float):
    # Transpose the batch (see https://stackoverflow.com/a/19343/3343043 for
    # detailed explanation). This converts batch-array of Transitions
    # to Transition of batch-arrays.
    batch = Transition(*zip(*transitions))

    # Compute a mask of non-final states and concatenate the batch elements
    # (a final state would've been the one after which simulation ended)
    non_final_mask = torch.tensor(tuple(map(lambda s: s is not None,
                                          batch.next_state)), device=device, dtype=torch.bool)
    non_final_next_states = torch.cat([s for s in batch.next_state
                                                if s is not None])
    state_batch = torch.cat(batch.state)
    action_batch = torch.cat(batch.action)
    reward_batch = torch.cat(batch.reward)

    # Compute Q(s_t, a) - the models computes Q(s_t), then we select the
    # columns of actions taken. These are the actions which would've been taken
    # for each batch state according to policy_net
    state_action_values = policy_net(state_batch).gather(1, action_batch)

    # Compute V(s_{t+1}) for all next states.
    # Expected values of actions for non_final_next_states are computed based
    # on the "older" target_net; selecting their best reward with max(1).values
    # This is merged based on the mask, such that we'll have either the expected
    # state value or 0 in case the state was final.
    next_state_values = torch.zeros(BATCH_SIZE, device=device)
    with torch.no_grad():
        # origin
        # next_state_values[non_final_mask] = target_net(non_final_next_states).max(1).values

        # double DQN
        next_state_values[non_final_mask] = target_net(non_final_next_states).gather(1, policy_net(non_final_next_states).argmax(dim=1, keepdim=True)).squeeze(1) # 降维度为n
    # Compute the expected Q values
    expected_state_action_values = (next_state_values * gamma) + reward_batch
    # Compute Huber loss
    elementwise_loss = F.smooth_l1_loss(input=state_action_values, target=expected_state_action_values.unsqueeze(1),
                                        reduction='none', beta=1.0)
    
    return elementwise_loss

def optimize_model():
    if len(memory) < BATCH_SIZE:
        return
    transitions, weights, indices = memory.sample(BATCH_SIZE, beta=BETA)
    weights = torch.FloatTensor(
        weights.reshape(-1, 1)
    ).to(device)
    
    elementwise_loss = _compute_dqn_loss(transitions=transitions, gamma=GAMMA)

    # PER: importance sampling before average
    loss = torch.mean(elementwise_loss * weights)

    # N-step Learning loss
    # we are gonna combine 1-step loss and n-step loss so as to
    # prevent high-variance. The original rainbow employs n-step loss only.
    if N_STEP > 1:
        gamma = GAMMA ** N_STEP
        samples = memory_n_step.sample_batch_from_idxs(indices)
        elementwise_loss_n_loss = _compute_dqn_loss(samples, gamma=gamma)
        elementwise_loss += elementwise_loss_n_loss
        
        # PER: importance sampling before average
        loss = torch.mean(elementwise_loss * weights)

    # Optimize the models
    optimizer.zero_grad()
    loss.backward()
    # In-place gradient clipping
    torch.nn.utils.clip_grad_value_(policy_net.parameters(), 100)
    optimizer.step()

    # PER: update priorities
    loss_for_prior = elementwise_loss.detach().cpu().numpy()
    new_priorities = loss_for_prior + PER_EPS
    memory.update_priorities(indices, new_priorities)

    return loss.item()


def save_best_model(path):
    try:
        #gpu tensor format save
        torch.save(obj={
            'policy_net_state_dict': policy_net.state_dict(),
        }, f=path)
        logger.info("save models success")
    except:
        logger.warning("save models failed.")


def decode_data(data):
    # load json
    recv_obser = json.loads(data)
    # print(recv_obser)
    observation = []
    observation_dict = recv_obser['observation']
    for key in observation_dict:
        observation.append(observation_dict[key])
    reward = recv_obser['reward']
    terminated = recv_obser['terminated']
    truncated = recv_obser['truncated']
    return observation, reward, terminated, truncated

def encode_data(init_delta, action, reset_flag=1):
    rudder = action - 5 + init_delta
 
    if rudder >= 30:
        rudder = 30
    elif rudder <= -30:
        rudder = -30
    
    origin_data = {'boatname':'SLM7001',
                   'restart': reset_flag,
                   'rudl': rudder,
                   'rudr': rudder,
                   'rspl': 1000,
                   'rspr': 1000,
                   'subSystem': "control"
                   }

    data = json.dumps(origin_data, sort_keys=True, indent=4, separators=(',', ':'))
    return data.encode('utf-8'), rudder


save_model_path = os.path.join(HOME, 'DQN/best_tracking_model.pt')
recorder_path = os.path.join(HOME, 'DQN/record.pkl')
# BATCH_SIZE is the number of transitions sampled from the replay buffer
# GAMMA is the discount factor as mentioned in the previous section
# EPS_START is the starting value of epsilon
# EPS_END is the final value of epsilon
# EPS_DECAY controls the rate of exponential decay of epsilon, higher means a slower decay
# TAU is the update rate of the target network
# LR is the learning rate of the ``AdamW`` optimizer
BATCH_SIZE = 128
GAMMA = 0.99
EPS_START = 0.9
EPS_END = 0.05
EPS_DECAY = 100_000
TAU = 0.005
LR = 6.25e-5
# MEMORY
BUFFER_SIZE = 100_000
# PER parameters
ALPHA = 0.5
BETA = 0.4
PER_EPS = 1e-5

# n-step
N_STEP = 3

# cell num
CELL_NUM=64


episode_durations = []
# Get number of actions from gym action space
n_actions = 11
action_space=range(n_actions) # 11个action
# Get the number of state observations
n_observations = 3
policy_net = DQN(n_observations, n_actions).to(device)
target_net = DQN(n_observations, n_actions).to(device)
target_net.load_state_dict(policy_net.state_dict())
optimizer = optim.AdamW(policy_net.parameters(), lr=LR, amsgrad=True)
# PER
memory = PrioritizedReplayBuffer(capacity=BUFFER_SIZE, alpha=ALPHA)
# memory for N-step Learning
if N_STEP > 1:
    memory_n_step = ReplayMemory(capacity=BUFFER_SIZE)

if __name__ == '__main__':
    # 与外部环境软件连接socket
    REMOTE_HOST = '127.0.0.1'
    REMOTE_PORT = 10901

    TCP_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    TCP_socket.connect((REMOTE_HOST, REMOTE_PORT))

    steps_done = 0
    num_episodes = 500
    frames_per_episode = 2000
    frame_idx = 0
    num_frames = num_episodes*frames_per_episode
    best_score = torch.tensor([0.0], dtype=torch.float32)

    # 训练中间数据记录 用于绘制图表

    # 训练效果记录
    recorder = dict({
        'frame': [],
        'epoch': [],
        'score': [],
        'loss': [],
        'epsilons': [],
    })
    
    # 初始舵角
    init_rudder = 0

    # 动作和状态记录 用于评价 在测试脚本中记录 不在训练脚本中记录
    for i_episode in range(num_episodes):
        action_data, init_rudder = encode_data(init_rudder, random.choice(action_space), reset_flag=1)
        # print(data)
        TCP_socket.send(action_data)
        # Initialize the environment and get its state
        info, addr = TCP_socket.recvfrom(1024)
        observation, reward, terminated, truncated= decode_data(info)
        state = observation
        state = torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0)
        score = float(0.0)
        while True:
            # PER: increase beta
            frame_idx += 1
            fraction = min(frame_idx / num_frames, 1.0)
            BETA = BETA + fraction * (1.0 - BETA)
            
            action, eps_threshold = select_action(state) # state
            action_data, init_rudder = encode_data(init_rudder,action.item(), reset_flag=0)
            TCP_socket.send(action_data)# action
            info, addr = TCP_socket.recvfrom(1024)
            observation, reward, terminated, truncated = decode_data(info)# next state
            
            reward = torch.tensor(reward, device=device).unsqueeze(0)# reward
            # print(f'result is {reward}')
            done = terminated or truncated
            score += reward
            if terminated:
                next_state = None
            else:
                next_state = torch.tensor(observation, dtype=torch.float32, device=device).unsqueeze(0)
            
            #   ('state', 'action','reward', 'next_state', 'done')
            transition = list()
            transition = [state, action, reward, next_state, done]

            # N-step transition
            if N_STEP:
                one_step_transition = memory_n_step.push(*transition)
            # 1-step transition
            else:
                one_step_transition = transition
                
            if one_step_transition:
                # Store the transition in memory
                memory.push(*one_step_transition)

            # Move to the next state
            state = next_state

            # Perform one step of the optimization (on the policy network)
            loss = optimize_model()

            # Soft update of the target network's weights
            # θ′ ← τ θ + (1 −τ )θ′
            target_net_state_dict = target_net.state_dict()
            policy_net_state_dict = policy_net.state_dict()
            for key in policy_net_state_dict:
                target_net_state_dict[key] = policy_net_state_dict[key]*TAU + target_net_state_dict[key]*(1-TAU)
            target_net.load_state_dict(target_net_state_dict)

            if done:
                init_rudder = 0
                
                if score.squeeze(0).cpu().numpy() >= best_score.squeeze(0).cpu().numpy():
                    best_score = score

                # save to recorder
                recorder['epoch'].append(i_episode)
                recorder['score'].append(score.squeeze(0).cpu().numpy())
                print(
                    f'episode is {i_episode}, score is {score.squeeze(0).cpu().numpy()}')
                break
            recorder["frame"].append(frame_idx)
            recorder['loss'].append(loss)
            recorder['epsilons'].append(eps_threshold)

    # save to file for ploting
    with open(recorder_path, 'wb') as f:
        pickle.dump(recorder, f)
        
    save_best_model(save_model_path)
    print('Complete')
